// 导入IntelliJ IDEA数据库工具的模型类。
import com.intellij.database.model.DasTable
import com.intellij.database.model.ObjectKind
// 导入用于大小写转换的工具类。
import com.intellij.database.util.Case
// 导入数据库工具的实用工具类。
import com.intellij.database.util.DasUtil
// 导入Java IO类。
import java.io.*
// 导入日期格式化工具类。
import java.text.SimpleDateFormat

// 定义上下文绑定的注释，这些是脚本可以访问的变量。
/*
 * Available context bindings:
 *   SELECTION   Iterable<DasObject> - 可迭代的数据库对象集合。
 *   PROJECT     project - 当前项目实例。
 *   FILES       files helper - 文件帮助工具。
 */

// 初始化包名变量。
packageName = ""
// 定义数据库类型到Java类型的映射关系。
typeMapping = [
        // 正则表达式匹配数据库类型，并映射到Java类型。
]

// 使用FILES工具提示用户选择保存生成文件的目录，并执行给定的闭包。
FILES.chooseDirectoryAndSave("Choose directory", "Choose where to store generated files") { dir ->
    // 过滤SELECTION中的DasTable对象，并且只选择类型为TABLE的表。
    SELECTION.filter { it instanceof DasTable && it.getKind() == ObjectKind.TABLE }.each { generate(it, dir) }
}

// define generate function that generates Java code for a given table
def generate(table, dir) {
    // Convert table name to a Java class name.
    def className = javaClassName(table.getName(), true)
    // Calculate fields for the table.
    def fields = calcFields(table)
    // Get the package name from the directory path.
    packageName = getPackageName(dir)
    // Create a PrintWriter to write to a Java file.
    PrintWriter printWriter = new PrintWriter(new OutputStreamWriter(new FileOutputStream(new File(dir, className + ".java")), "UTF-8"))
    // Use the PrintWriter to generate the Java class.
    printWriter.withPrintWriter { out -> generate(out, className, fields, table) }
}

// Define getPackageName function to convert a directory path to a package name.
def getPackageName(dir) {
    // Replace backslashes and slashes with dots, and remove the 'src' part.
    return dir.toString().replaceAll("\\\\", ".").replaceAll("/", ".").replaceAll("^.*src(\\.main\\.java\\.)?", "") + ";"
}

// Define generate function to write Java class code to the PrintWriter.
def generate(out, className, fields, table) {
    // Print package declaration.
    out.println "package $packageName"
    // Print imports for MyBatis Plus annotations and Lombok annotations.
    out.println ""
    out.println "import com.baomidou.mybatisplus.annotation.TableName;"
    out.println "import com.baomidou.mybatisplus.annotation.TableId;"
    out.println "import com.baomidou.mybatisplus.annotation.IdType;"
    out.println "import com.baomidou.mybatisplus.annotation.TableField;"
    out.println "import java.io.Serializable;"
    out.println "import lombok.*;"
    // Collect all unique types used in the fields.
    Set types = new HashSet()
    fields.each() {
        types.add(it.type)
    }
    // If any field is of type Date or InputStream, import the respective class.
    if (types.contains("Date")) {
        out.println "import java.util.Date;"
    }
    if (types.contains("InputStream")) {
        out.println "import java.io.InputStream;"
    }
    // Print class comment.
    out.println ""
    out.println "/**"
    out.println " * @Description"
    out.println " * @Author  wwh"
    // Print the current date in the comment.
    out.println " * @Date " + new SimpleDateFormat("yyyy-MM-dd").format(new Date())
    out.println " */"
    // Print Lombok annotations for data class, all-args constructor, no-args constructor, and table name.
    out.println ""
    out.println "@Data"
    out.println "@AllArgsConstructor"
    out.println "@NoArgsConstructor"
    out.println "@TableName(\"${table.getName()}\")"
    // Print class declaration.
    out.println "public class $className implements Serializable {"
    // Print serialVersionUID.
    out.println ""
    out.println genSerialID()
    // Print field definitions with annotations and comments.
    fields.each() {
        out.println ""
        if (isNotEmpty(it.comment)) {
            out.println "\t/**"
            out.println "\t * ${it.comment.toString()}"
            out.println "\t */"
        }
        if (it.annos != "") out.println "   ${it.annos}"
        out.println "\tprivate ${it.type} ${it.name};"
    }
    // Close the class declaration.
    out.println ""
    out.println "}"
}

// Define calcFields function to calculate field definitions for a table.
def calcFields(table) {
    // Use DasUtil.getColumns to get columns and reduce them to a list of field definitions.
    DasUtil.getColumns(table).reduce([]) { fields, col ->
        // Apply lower case to the column data type specification.
        def spec = Case.LOWER.apply(col.getDataType().getSpecification())
        // Find the matching Java type from the type mapping.
        def typeStr = typeMapping.find { p, t -> p.matcher(spec).find() }.value
        // Create a map of field properties.
        def comm = [
                colName: col.getName(),
                name: javaName(col.getName(), false),
                type: typeStr,
                comment: col.getComment(),
                annos: DasUtil.isPrimary(col) ? "\t@TableId(type = IdType.AUTO)" : "\t@TableField(\"${col.getName()}\")"
        ]
        fields += [comm]
    }
}

// Define javaClassName function to convert a string to a Java class name.
def javaClassName(str, capitalize) {
    // Split the string into words, capitalize them, and join them.
    def s = com.intellij.psi.codeStyle.NameUtil.splitNameIntoWords(str)
            .collect { Case.LOWER.apply(it).capitalize() }
            .join("")
    // Replace non-Java identifier characters with underscores.
            .replaceAll(/[^\p{javaJavaIdentifierPart}[_]]/, "_")
    // Capitalize the first letter if needed.
    capitalize || s.length() == 1 ? s : Case.LOWER.apply(s[0]) + s[1..-1]
}

// Define javaName function similar to javaClassName but does not force capitalization.
def javaName(str, capitalize) {
    // Similar to javaClassName but without forcing the first letter to be capitalized.
    def s = com.intellij.psi.codeStyle.NameUtil.splitNameIntoWords(str)
            .collect { Case.LOWER.apply(it).capitalize() }
            .join("")
            .replaceAll(/[^\p{javaJavaIdentifierPart}[_]]/, "_")
    capitalize || s.length() == 1 ? s : Case.LOWER.apply(s[0]) + s[1..-1]
}

// Define isNotEmpty function to check if a content is not null and not just whitespace.
def isNotEmpty(content) {
    return content != null && content.toString().trim().length() > 0
}

// Define genSerialID function to generate a serial version UID for the class.
static String genSerialID() {
    return "\tprivate static final long serialVersionUID = " + Math.abs(new Random().nextLong()) + "L;"
}